import numpy as np
import matplotlib.pyplot as plt
from init_control import UAVParameters
from uav_dynamics import UAVDynamics
from common_function import rotation_matrix
from PIDController import posion_pos_controller
from LQRController import LQRController
from square_signal import generate_ramp_signal
# from self_LQRController import LQRController
from FIS_LQRController import LQRController

class UAVSimulation:
    def __init__(self, params, dynamics, position_controller, velocity_controller, pose_controllers):
        self.params = params
        self.dynamics = dynamics
        self.position_history = []
        self.setpoint_history = []
        self.position_controller = position_controller
        self.velocity_controller = velocity_controller
        self.roll_controller = pose_controllers[0]
        self.pitch_controller = pose_controllers[1]
        self.yaw_controller = pose_controllers[2]

    def motor_mixer(self, roll, pitch, yaw, thrust):
        idle_PWM = 1000
        scale = 1000
        M1 = (thrust - roll + pitch + yaw) * scale + idle_PWM
        M2 = (thrust + roll - pitch + yaw) * scale + idle_PWM
        M3 = (thrust + roll + pitch - yaw) * scale + idle_PWM
        M4 = (thrust - roll - pitch - yaw) * scale + idle_PWM
        return M1, M2, M3, M4

    def desired_euler_angles(self, eax, eay, psi):
        g = 9.8  # 重力加速度 (m/s^2)
        phi = (-np.sin(psi) * eax + np.cos(psi) * eay) / g
        theta = (-np.cos(psi) * eax - np.sin(psi) * eay) / g
        return phi, theta

    def pwm_to_rpm(self, pwm):
        min_pwm = 1000
        max_pwm = 2000
        pwm = np.clip(pwm, min_pwm, max_pwm)
        rpm = np.clip((pwm - min_pwm) / (max_pwm - min_pwm), -1, 1)
        return rpm

    def run(self, square_corners, duration, dt):
        DEG2RAD = self.params.DEG2RAD
        posE = self.params.ModelInit_PosE.astype(np.float64)
        velB = self.params.ModelInit_VelB.astype(np.float64)
        velE = self.params.ModelInit_VelB.astype(np.float64)
        angEuler = self.params.ModelInit_AngEuler
        rate = self.params.ModelInit_RateB
        ModelParam_motorCr = self.params.ModelParam_motorCr
        ModelParam_motorWb = self.params.ModelParam_motorWb
        self.position_history.append(posE)

        state = [velB, angEuler, rate]
        time_steps = int(duration / dt)
        num_steps = time_steps // (len(square_corners) - 1)
        input_signal = generate_ramp_signal(square_corners, num_steps)
        vel_setpoint = np.array([0, 0, 0])
        acc_setpoint = np.array([0, 0, 0])
        velE_dot = np.array([0, 0, 0])

        for step in range(time_steps):
            # setpoint = input_signal[step]
            setpoint = square_corners[step//num_steps]
            # 位置环LQR控制
            vel_setpoint = self.position_controller.update(setpoint,posE,vel_setpoint-velE)

            # 速度环LQR控制
            acc_setpoint = self.velocity_controller.update(vel_setpoint,velE,acc_setpoint-velE_dot)

            # 姿态环LQR控制
            desired_phi, desired_theta = self.desired_euler_angles(acc_setpoint[0], acc_setpoint[1], angEuler[2])
            phi = np.clip(desired_phi, -35 * DEG2RAD, 35 * DEG2RAD)
            theta = np.clip(desired_theta, -35 * DEG2RAD, 35 * DEG2RAD)

            roll_output = np.clip(self.roll_controller.update(outer_setpoint=phi, outer_measurement=angEuler[0],
                                                              inner_measurement=rate[0], dt=dt), -1, 1)
            pitch_output = np.clip(self.pitch_controller.update(outer_setpoint=theta, outer_measurement=angEuler[1],
                                                                inner_measurement=rate[1], dt=dt), -1, 1)
            yaw_output = np.clip(self.yaw_controller.update(outer_setpoint=0, outer_measurement=angEuler[2],
                                                            inner_measurement=rate[2], dt=dt), -1, 1)
            thrust = np.clip(-acc_setpoint[2] + self.params.THR_HOVER, 0.05, 0.9)

            M1, M2, M3, M4 = self.motor_mixer(roll_output, pitch_output, yaw_output, thrust)

            w_target = (np.array(
                [self.pwm_to_rpm(M1), self.pwm_to_rpm(M2), self.pwm_to_rpm(M3),
                 self.pwm_to_rpm(M4)]) * ModelParam_motorCr +
                        np.array([ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb, ModelParam_motorWb]))

            state,velB_dot = self.dynamics.update_dynamics(state, w_target, dt)
            velB, angEuler, rate = state
            R = rotation_matrix(angEuler)
            velE_dot = R@velB_dot
            velE = R @ velB
            posE = posE + velE * dt
            self.position_history.append(posE.copy())
            print(
                f"Time: {step * dt:.2f}s, Position: {posE}, Velocity_world: {velE}, Angle: {angEuler}, Rate: {rate}")
            self.setpoint_history.append(setpoint.copy())

    def plot_trajectory(self):
        position_history = np.array(self.position_history)
        setpoint_history = np.array(self.setpoint_history)
        fig = plt.figure()
        ax = fig.add_subplot(121, projection='3d')
        ax.plot(position_history[:, 0], position_history[:, 1], -position_history[:, 2], label='UAV Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title('Trajectory of UAV')
        ax.legend()
        ax = fig.add_subplot(122, projection='3d')
        ax.plot(setpoint_history[:, 0], setpoint_history[:, 1], -setpoint_history[:, 2], label='Setpoint Trajectory')
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')
        ax.set_title('Desired Trajectory')
        ax.legend()

    def plot_state_error(self):
        self.position_controller.plot_error()
        self.velocity_controller.plot_error()
        self.roll_controller.outer_controller.plot_error('roll_angle')
        self.roll_controller.inner_controller.plot_error('roll_rate')
        self.pitch_controller.outer_controller.plot_error('pitch_angle')
        self.pitch_controller.inner_controller.plot_error('pitch_rate')
        self.yaw_controller.outer_controller.plot_error('yaw_angle')
        self.yaw_controller.inner_controller.plot_error('yaw_rate')


if __name__ == '__main__':
    params = UAVParameters()
    dynamics = UAVDynamics(params)

    # 定义位置环的状态空间模型
    A_pos = np.zeros((3, 3))
    B_pos = np.eye(3)

    Q_pos = np.eye(3) * 0.8
    R_pos = np.eye(3) * 2.0

    # 定义速度环的状态空间模型
    A_vel = np.zeros((3, 3))
    B_vel = np.eye(3)

    Q_vel = np.eye(3) * 4
    R_vel = np.eye(3) * 1.2

    # 定义姿态环的状态空间模型
    A_att = np.zeros((6, 6))
    A_att[0, 3] = 1
    A_att[1, 4] = 1
    A_att[2, 5] = 1

    B_att = np.zeros((6, 3))
    B_att[3, 0] = 1
    B_att[4, 1] = 1
    B_att[5, 2] = 1

    Q_att = np.eye(6) * 0.05
    R_att = np.eye(3) * 2

    # 创建LQR控制器
    position_controller = LQRController(A_pos, B_pos, Q_pos, R_pos)
    velocity_controller = LQRController(A_vel, B_vel, Q_vel, R_vel)
    # position_controller = AdaptiveLQRController(A_pos, B_pos, Q_pos, R_pos)
    # velocity_controller = AdaptiveLQRController(A_vel, B_vel, Q_vel, R_vel)
    pose_controller = posion_pos_controller(params).pos_controller()

    # 目标轨迹
    target_altitude = -5
    square_corners = [
        params.ModelInit_PosE.astype(np.float64),
        np.array([0, 0, target_altitude]),
        np.array([5, 5, target_altitude]),
        np.array([5, 0, target_altitude]),
        np.array([0, 5, target_altitude]),
        np.array([0, 0, target_altitude]),
        np.array([0, 0, 0]),
        # np.array([0, 0, 0])
    ]

    # 运行仿真
    simulation = UAVSimulation(params, dynamics, position_controller, velocity_controller, pose_controller)
    simulation.run(square_corners, 70.0, 0.01)
    simulation.plot_trajectory()
    simulation.plot_state_error()
    plt.show()
